#include "Image3DTexture.h"
#include "../../Core/CoreGlobals.h"
#include "../../Utils/LogUtil.h"
#include "../../Utils/FileUtil.h"
#include "../../Templates/NewObject.h"
#include "../RenderState.h"
#include "../../Game/Process/RenderProcess.h"

Image3DTexture::Builder::Builder() noexcept = default;
Image3DTexture::Builder::~Builder() noexcept = default;
Image3DTexture::Builder::Builder(Image3DTexture::Builder const& rhs) noexcept = default;
Image3DTexture::Builder::Builder(Image3DTexture::Builder && rhs) noexcept = default;
Image3DTexture::Builder& Image3DTexture::Builder::operator=(Image3DTexture::Builder const& rhs) noexcept = default;
Image3DTexture::Builder& Image3DTexture::Builder::operator=(Image3DTexture::Builder && rhs) noexcept = default;

Image3DTexture::Builder& Image3DTexture::Builder::SetMultiple(bool isMultiple) noexcept
{
	mImpl->isMultiple = isMultiple;
	return *this;
}

Image3DTexture::Builder& Image3DTexture::Builder::SetFormat(VkFormat format) noexcept
{
	mImpl->format = format;
	return *this;
}

Image3DTexture::Builder& Image3DTexture::Builder::SetExtent(VkExtent3D extent) noexcept
{
	mImpl->extent = extent;
	return *this;
}

Image3DTexture::Builder& Image3DTexture::Builder::SetWidth(uint32 width) noexcept
{
	mImpl->extent.width = width;
	return *this;
}

Image3DTexture::Builder& Image3DTexture::Builder::SetHeight(uint32 height) noexcept
{
	mImpl->extent.height = height;
	return *this;
}

Image3DTexture::Builder& Image3DTexture::Builder::SetDepth(uint32 depth) noexcept
{
	mImpl->extent.depth = depth;
	return *this;
}

Image3DTexture::Builder& Image3DTexture::Builder::SetSharingMode(VkSharingMode sharingMode) noexcept
{
	mImpl->sharingMode = sharingMode;
	return *this;
}

Image3DTexture::Builder& Image3DTexture::Builder::AddQueueFamily(VkQueueFlagBits queueFlag) noexcept
{
	mImpl->queueFamilys.push_back(queueFlag);
	return *this;
}


Image3DTexture::Builder& Image3DTexture::Builder::SetData(void* data) noexcept
{
	mImpl->data = data;
	return *this;
}

RefCountPtr<Image3DTexture> Image3DTexture::Builder::Build()
{
	assert(IsInRenderThread());
	return RefCountPtr<Image3DTexture>(NewObject<Image3DTexture>(*mImpl));
}

RefCountPtr<Image3DTexture> Image3DTexture::Builder::Build(RenderProcess* process, const std::string& name)
{
	assert(IsInRenderThread());
	Image3DTexture* image3DTexture = NewObject<Image3DTexture>(*mImpl);
	process->AddTexture(name, image3DTexture);
	return RefCountPtr<Image3DTexture>(image3DTexture);
}

void Image3DTexture::ReleaseResource()
{
	if (IsInRenderThread())
	{
		if (mProcess != nullptr)
		{
			mProcess->RemoveTexture(mName, this);
		}
		DeleteObject<Image3DTexture>(this);
	}
	else
	{
		AsyncTask(GRenderThreadId, [this]() {
			if (mProcess != nullptr)
			{
				mProcess->RemoveTexture(mName, this);
			}
			DeleteObject<Image3DTexture>(this);
			});
	}
}

Image3DTexture::Image3DTexture(BuilderDetails & builderDetails)
{
	mIsMultiple = builderDetails.isMultiple;
	mIsRenderTarget = false;

	mFlags = 0;
	mImageType = VK_IMAGE_TYPE_3D;
	mFormat = builderDetails.format;
	mExtent = builderDetails.extent;
	mUsage = VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_SAMPLED_BIT;
	mMipLevels = 1;
	mArrayLayers = 1;
	mSamples = VK_SAMPLE_COUNT_1_BIT;
	mTiling = VK_IMAGE_TILING_OPTIMAL;
	mSharingMode = builderDetails.sharingMode;
	mQueueFamilys.assign(builderDetails.queueFamilys.begin(), builderDetails.queueFamilys.end());
	if (mQueueFamilys.size() > 1)
	{
		// Check if the queueFamilys is same
		bool isDifferent = false;
		const uint32 standQueueFamily = GVulkanRHI->GetQueueFamily(mQueueFamilys[0]);
		for (uint32 i = 1; i < mQueueFamilys.size(); ++i)
		{
			if (GVulkanRHI->GetQueueFamily(mQueueFamilys[i]) != standQueueFamily)
				isDifferent = true;
		}

		if (isDifferent)
			mSharingMode = VK_SHARING_MODE_CONCURRENT; // multi queue
		else
			mSharingMode = VK_SHARING_MODE_EXCLUSIVE;
	}
	else
	{
		assert(mSharingMode == VK_SHARING_MODE_EXCLUSIVE);
	}
	mImageLayouts.resize(GetFrameCount(), VK_IMAGE_LAYOUT_UNDEFINED);
	mMemoryProperty = VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT;

	mViewType = VK_IMAGE_VIEW_TYPE_3D;
	mComponents = { VK_COMPONENT_SWIZZLE_R , VK_COMPONENT_SWIZZLE_G , VK_COMPONENT_SWIZZLE_B , VK_COMPONENT_SWIZZLE_A };
	// No mipmap
	mSubresourceRange = { VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0 , 1 };

	mImages.resize(GetFrameCount(), VK_NULL_HANDLE);
	mImageMemorys.resize(GetFrameCount(), VK_NULL_HANDLE);
	mImageViews.resize(GetFrameCount(), VK_NULL_HANDLE);

	for (uint32 i = 0; i < GetFrameCount(); ++i)
	{
		GVulkanRHI->CreateImage(mImages[i], mImageMemorys[i], mFlags, mExtent, mUsage, mFormat, mMipLevels, mMemoryProperty, mSamples, mImageType, mImageLayouts[i], mArrayLayers, mTiling, mSharingMode, mQueueFamilys);
	}

	// if the data is not null, create the buffer to copy data to image
	if (builderDetails.data != nullptr)
	{

		VkBuffer stagingBuffer;
		VkDeviceMemory stagingMemory;

		VkDeviceSize dataSize = static_cast<VkDeviceSize>(mExtent.width * mExtent.height * mExtent.depth * GVkFormatSize[mFormat]);

		GVulkanRHI->CreateBuffer(stagingBuffer, stagingMemory, dataSize, VK_BUFFER_USAGE_TRANSFER_SRC_BIT, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT);

		void* data;
		VK_CHECK_RESULT(vkMapMemory(GVulkanRHI->GetLogicalDevice(), stagingMemory, 0, dataSize, 0, &data));
		memcpy(data, builderDetails.data, dataSize);
		vkUnmapMemory(GVulkanRHI->GetLogicalDevice(), stagingMemory);

		VkCommandBuffer commandBuffer = GRenderState->BeginSingleCommand();

		for (uint32 i = 0; i < GetFrameCount(); ++i)
		{
			// TransitionImageLayout for copy data
			GVulkanRHI->TransferLayout(commandBuffer, mImages[i], mImageLayouts[i], VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, mSubresourceRange, VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT);
			mImageLayouts[i] = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL;

			GVulkanRHI->CopyBufferToImage(commandBuffer, stagingBuffer, mImages[i], mExtent, { VK_IMAGE_ASPECT_COLOR_BIT, 0, 0, mSubresourceRange.layerCount });

			GVulkanRHI->TransferLayout(commandBuffer, mImages[i], mImageLayouts[i], VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL, mSubresourceRange, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT);
			mImageLayouts[i] = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
		}

		GRenderState->FlushSingleCommand(commandBuffer);

		vkDestroyBuffer(GVulkanRHI->GetLogicalDevice(), stagingBuffer, nullptr);
		vkFreeMemory(GVulkanRHI->GetLogicalDevice(), stagingMemory, nullptr);
	}

	for (uint32 i = 0; i < GetFrameCount(); ++i)
	{
		GVulkanRHI->CreateImageView(mImageViews[i], mImages[i], mFormat, mViewType, mSubresourceRange, mComponents);
	}
}
